import numpy as np
import pandas as pd

from tools.change_cwd_to_main import change_cwd_to_main

change_cwd_to_main()


from predict_gold_price import *
from tools.visualize_data import tabular_output

flag__training = 1
# flag__training = 0
if flag__training:
    tt.tqdm_sleep('===== Start Train! =====')
    run_train()
    tt.tqdm_sleep('===== Start Evaluation! =====')

use_best_model = True
# use_best_model = False
best_model_path = f'tempdir/models/{experiment_name}/best.pth'
newest_model_path = f'tempdir/models/{experiment_name}/newest.pth'
# model.load_state_dict()
model.load_state_dict(torch.load(best_model_path)) if use_best_model else model.load_state_dict(torch.load(newest_model_path))
# model.eval()
device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
# device
model = model.to(device)


# train((X_train, X_test, y_train, y_test), model, optimizer, criterion,
#       total_epoch=total_epoch, log_interval=log_interval,
#       experiment_name=experiment_name, save=True, is_classification=is_classification, scheduler=scheduler, scheduler_step_size=scheduler_step_size)


print('\n\n======================= on test ========================\n')

print(f'====== baseline_dc:')
show_json(baseline_dc)

test_size
window_size
# X_test[-1]
day_i = -1  # 14号截止
test_i: torch.FloatTensor = X_test[day_i:, :, :]
# test_i.unsqueeze_(0)
test_i.shape
X_test.shape

test_i = test_i.to(device)
out: torch.FloatTensor = model(test_i)
out.shape
out
out = out.cpu().detach().numpy()
print(f'--- out: {out}')

predict_scaler = scaler_dc.get(predict_column)

res = inverse_transform(predict_scaler, out)
print(f'--- res: {res[0][0] * 100: .3f} %')


# --- pred train x
X_train = X_train.to(device)
train_y_pred = model(X_train)
train_y_pred = train_y_pred.cpu().detach().numpy()
train_y_pred = inverse_transform(predict_scaler, train_y_pred * 100)
train_y_true = inverse_transform(predict_scaler, y_train.reshape(-1, 1) * 100)

# print(tabular_output(np.vstack((np.arange(train_y_pred.shape[-1]), train_y_pred, train_y_true)), decimal_places=2))
train_y_pred.shape, train_y_true.shape
np.hstack((train_y_pred, train_y_true))
df_train_pred = pd.DataFrame(np.hstack((train_y_true, train_y_pred)), columns=['true', 'pred'])
max_visualize_rows = 1000
if len(df_train_pred) > max_visualize_rows:
      df_train_pred = df_train_pred.sample(n=max_visualize_rows, random_state=42)
visualize_data(df_train_pred, title="Train: compare true and pred")


X_test = X_test.to(device)

# X_test = X_train.to(device)
# model(X_test)
# y_test.shape
# y_test.reshape(1, -1).shape

_y_pred = model(X_test)
y_pred = _y_pred.cpu().detach()
# np.round(y_pred.numpy().tolist(), 3).tolist()
y_pred.shape
# y_pred = y_pred[0] if isinstance(y_pred, tuple) else y_pred

y_pred.resize_as_(y_test)
eval_res = test_by_evaluate_metrics(y_pred, y_true)
print(f'--- eval_res: {eval_res}')
# torch.mean(y_true)
# torch.mean(y_train)
# torch.mean(y_pred)

torch.round((y_pred - torch.mean(y_pred)) * 1000) / 1000
# (y_pred - torch.mean(y_pred)).round(3)
# X_test

# test_by_evaluate_metrics(np.array([1, 1, 1]), np.array([0.8, 1, -1]))


line_pred = np.round(inverse_transform(predict_scaler, y_pred.reshape(1, -1)) * 100, 3)
line_test = np.round(inverse_transform(predict_scaler, y_test.reshape(1, -1)) * 100, 3)

from sklearn.metrics import mean_absolute_error, mean_squared_error
mean_absolute_error(y_pred, y_test)
mean_squared_error(y_pred, y_test)

print(f"--- real mape: {mean_absolute_percentage_error(line_pred, line_test): .4f},"
      f" mae: {mean_absolute_error(line_pred, line_test): .4f},"
      f" mse: {mean_squared_error(line_pred, line_test)}: .4f")


from bdtime import show_ls

print('------ compare pred and true ------')
show_ls(np.round(np.vstack((line_pred, line_test)), 4).tolist())
print('------ compare pred and true ------')

# 调用函数并打印结果
print(tabular_output(np.vstack((np.arange(line_pred.shape[-1]), line_pred, line_test)), decimal_places=2))

# --- 可视化
print(f'--- y_train mean: {np.mean(y_train.numpy()): .4f},'
      f' y_test mean: {np.mean(y_test.numpy()): .4f},'
      f' _y_pred mean: {np.mean(y_pred.numpy()): .4f}')
print(f'--- line_pred mean: {np.mean(line_pred): .4f}, line_test mean: {np.mean(line_test): .4f}')
print(f'------ describe_data(line_pred): {describe_data(line_pred[0])}')
print()

df_1 = pd.DataFrame(np.vstack((line_test, line_pred)).T, columns=['true', 'pred'])
visualize_data(df_1, title="Test: compare true and pred")


# show_data_density_by_kde(line_pred[0], 'line_pred')
# show_data_density_by_kde(line_test[0], 'line_test')

# np.sum(line_pred)
# np.sum(line_test)


# future_pred_days = 10
# for i in range(future_pred_days):
#     y_pred = model.predict(X_test[-1, :, :].reshape(1, -1, 1))


# def eval():
#     model.eval()
#     for i in range(fut_pred):
#         seq = torch.FloatTensor(test_inputs[-train_window:])
